
#include "websocketproxy.h"


WebSocketProxy::WebSocketProxy(const QString& destination, int destPort, QObject *parent) : QObject(parent)
{
    p_server = new WebSocketServer(this);
    v_destinationPort = destPort;
    v_destinationHostname = destination;
}


void WebSocketProxy::listen(int port)
{
    p_server->listen(QHostAddress::Any, port);
    QObject::connect(p_server, SIGNAL(socketReady()), this, SLOT(newClient1Connected()));
}

void WebSocketProxy::newClient1Connected()
{
    WebSocket *newClient1 = p_server->nextPendingConnection();
    Socket *newClient2 = new Socket();
    v_waitingClient.insert(newClient2, newClient1);
    QObject::connect(newClient2, SIGNAL(connected()), this, SLOT(newClient2Connected()));

    QObject::connect(newClient2, SIGNAL(error(QAbstractSocket::SocketError)), this, SLOT(processClient2Error(QAbstractSocket::SocketError)));
    
    QObject::connect(newClient1, SIGNAL(error(QAbstractSocket::SocketError)), this, SLOT(processClient1Error(QAbstractSocket::SocketError)));


    newClient2->connectToHost(v_destinationHostname, v_destinationPort);
    qDebug() << "new incoming websocket connection";
}


void WebSocketProxy::newClient2Connected()
{
    Socket *newClient2 = static_cast<Socket*>(sender());
    WebSocket *newClient1 = v_waitingClient.take(newClient2);
    v_client1ToClient2.insert(newClient1, newClient2);
    v_client2ToClient1.insert(newClient2, newClient1);

    QObject::connect(newClient1, SIGNAL(dataReceived(QByteArray)), this, SLOT(dataReceivedFromClient1(QByteArray)));
    QObject::connect(newClient2, SIGNAL(dataReceived(QByteArray)), this, SLOT(dataReceivedFromClient2(QByteArray)));

    QObject::connect(newClient1, SIGNAL(disconnected()), this, SLOT(disconnectClient2()));
    QObject::connect(newClient2, SIGNAL(disconnected()), this, SLOT(disconnectClient1()));
    

    qDebug() << "destination socket connected";
    Q_EMIT newTunnel();
}

void WebSocketProxy::processClient1Error(QAbstractSocket::SocketError error)
{
    qDebug() << "Client 1 error "<< error;
    WebSocket *client1 = static_cast<WebSocket*>(sender());
    Socket *client2 = NULL;
    client2 = removeLinkClient1ToClient2(client1);
    client2->close();
}

void WebSocketProxy::processClient2Error(QAbstractSocket::SocketError error)
{
    qDebug() << "Client 2 error "<< error;
    Socket *client2 = static_cast<Socket*>(sender());
    WebSocket *client1 = NULL;
    if((client1 = v_waitingClient.take(client2)) == 0){
        client1 = removeLinkClient2ToClient1(client2);
    }
    client1->close();
}

void WebSocketProxy::disconnectClient1()
{
    qDebug() << "disconnect client 1";
    Socket *client2 = static_cast<Socket*>(sender());
    WebSocket *client1 = NULL;
    client1 = removeLinkClient2ToClient1(client2);
    client1->disconnect();
    client1->close();

    client1->deleteLater();
    client2->deleteLater();
}

void WebSocketProxy::disconnectClient2()
{
    qDebug() << "disconnect client 2";
    WebSocket *client1 = static_cast<WebSocket*>(sender());
    Socket *client2 = NULL;
    client2 = removeLinkClient1ToClient2(client1);
    client2->disconnect();
    client2->close();
    
    client1->deleteLater();
    client2->deleteLater();
}


void WebSocketProxy::dataReceivedFromClient1(QByteArray data)
{
    WebSocket *client1 = static_cast<WebSocket*>(sender());
    Socket *client2 = v_client1ToClient2.value(client1);
    client2->write(data);
}

void WebSocketProxy::dataReceivedFromClient2(QByteArray data)
{
    Socket *client2 = static_cast<Socket*>(sender());
    WebSocket *client1 = v_client2ToClient1.value(client2);
    client1->write(data);
}


Socket* WebSocketProxy::removeLinkClient1ToClient2(WebSocket* client1)
{
    Socket *client2 = v_client1ToClient2.take(client1);
    v_client2ToClient1.remove(client2);
    return client2;
}

WebSocket* WebSocketProxy::removeLinkClient2ToClient1(Socket* client2)
{
    WebSocket *client1 = v_client2ToClient1.take(client2);
    v_client1ToClient2.remove(client1);
    return client1;
}

